#python3 ocr.py -i 1.jpg -o 1.txt
from paddleocr import PaddleOCR, draw_ocr
import sys,datetime
import getopt
import os,fitz
import glob
from PIL import Image
pdf_dir = []

def exe_ocr(img_path,file_txt):
    ocr = PaddleOCR(use_angle_cls=False, lang="ch")  
    result = ocr.ocr(img_path, cls=False)
    res = result[0]    
    boxes = []    
    txt = ""    

    for line in res:
        txt += line[1][0]+"\n"    
        boxes.append(line[0])    

    with open(file_txt, 'a')as f:    
        f.write(txt)

 
def get_file():
    docunames = os.listdir()
    for docuname in docunames:
        if os.path.splitext(docuname)[1] == '.pdf':  
            pdf_dir.append(docuname)
 
 
def conver_img():
    for pdf in pdf_dir:
        doc = fitz.open(pdf)
        pdf_name = os.path.splitext(pdf)[0]
        for pg in range(doc.page_count):
            page = doc[pg]
            rotate = int(0)
            zoom_x = 3.0
            zoom_y = 3.0
            trans = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
            pm = page.get_pixmap(matrix=trans, alpha=False)
            pm.save('%s.png' % pg)
 
def main(): 
    now = datetime.datetime.now()
    file_txt =now.strftime("%Y-%m-%d_%H:%M:%S") + ".txt" 
    folder_path = os.path.dirname(os.path.abspath(__file__))
    files = os.listdir(folder_path)
    files = glob.glob(folder_path + '/*')
    for file in files:
        if os.path.splitext(file)[-1]=='.png':
        # img_result = file_txt[:file_txt.rindex('.')+1]+"jpg"
            exe_ocr(file,file_txt)

def del_png():
# 获取当前工作目录
    current_dir = os.getcwd() 
    # 遍历当前目录中所有文件和子目录
    for root, dirs, files in os.walk(current_dir):
        for file in files:
            # 判断文件名后缀为指定类型（这里示范删除txt文件）
            if file.endswith(".png"):
                # 构建完整路径
                filepath = os.path.join(root, file)
                
                try:
                    # 删除文件
                    os.remove(filepath)
                    # print("成功删除文件：", filepath)
                except OSError as e:
                    print("无法删除文件：", filepath, "错误信息：", str(e))

if __name__ == '__main__':
    #1.读取PDF文件
    get_file()
    #2。转成PNG图片
    conver_img()
    #3文字识别
    main()
    del_png()